package vip.wangwenhao.autoconfigure.aspect;

import com.google.common.collect.ImmutableList;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import vip.wangwenhao.autoconfigure.service.CacheKeyGenerator;
import vip.wangwenhao.common.annotations.LimitLock;
import vip.wangwenhao.common.enums.LimitType;
import vip.wangwenhao.common.exception.LimitException;
import vip.wangwenhao.common.util.ExceptionUtils;
import vip.wangwenhao.common.util.WebUtils;

import java.io.Serializable;
import java.lang.reflect.Method;

/**
 * 限流增强类 - 基于redis分布式锁
 *
 * @author wangwenhao
 */
@Aspect
public class LimitAspect {

    @Autowired(required = false)
    private RedisTemplate<String, Serializable> limitRedisTemplate;
    @Autowired
    private CacheKeyGenerator cacheKeyGenerator;

    @Around("execution(public * *(..)) && @within(vip.wangwenhao.common.annotations.LimitLock)")
    public Object interceptor(ProceedingJoinPoint pjp) {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        LimitLock limitLockAnnotation = method.getAnnotation(LimitLock.class);
        LimitType limitType = limitLockAnnotation.limitType();
        String key = null;
        int limitPeriod = limitLockAnnotation.period();
        int limitCount = limitLockAnnotation.count();
        switch (limitType) {
            case IP:
                key = WebUtils.getRemoteIp();
                break;
            case CUSTOMER:
                //此处想根据表达式或者一些规则生成
                key = cacheKeyGenerator.getLockKey(pjp);
                break;
            case METHOD:
                String serviceName = pjp.getTarget().getClass().getName();
                key = serviceName + StringUtils.upperCase(method.getName());

            default:
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitLockAnnotation.prefix(), key));
        try {
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
            if (count != null && count.intValue() <= limitCount) {
                return pjp.proceed();
            } else {
                throw new LimitException();
            }
        } catch (Throwable e) {
            throw ExceptionUtils.exception(e);
        }
    }

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    public String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

}